"""
翁氏宗谱RAG系统 - 主程序入口

这是RAG系统的主程序入口，提供多种运行模式：
1. 测试模式：运行基础功能测试
2. 初始化模式：初始化系统和处理PDF文件
3. 问答模式：启动问答服务
4. Web模式：启动Web界面
5. 第二阶段测试：运行第二阶段功能测试

使用方法：
python main.py --mode test    # 运行测试
python main.py --mode init    # 初始化系统
python main.py --mode qa      # 问答模式
python main.py --mode web     # Web界面模式
python main.py --mode stage2  # 第二阶段测试

作者：xiaodeme@163.com
版本：1.0.0
"""

import argparse
import sys
import os
import logging

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

# 导入配置模块
try:
    from src.config import config
except ImportError:
    from config import config

# 配置日志
logging.basicConfig(
    level=getattr(logging, config.LOG_LEVEL),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def print_banner():
    """打印程序横幅"""
    banner = """
    ████████╗███╗   ██╗███████╗    ██████╗  █████╗  ██████╗ 
    ╚══██╔══╝████╗  ██║██╔════╝    ██╔══██╗██╔══██╗██╔════╝ 
       ██║   ██╔██╗ ██║███████╗    ██████╔╝███████║██║  ███╗
       ██║   ██║╚██╗██║╚════██║    ██╔══██╗██╔══██║██║   ██║
       ██║   ██║ ╚████║███████║    ██║  ██║██║  ██║╚██████╔╝
       ╚═╝   ╚═╝  ╚═══╝╚══════╝    ╚═╝  ╚═╝╚═╝  ╚═╝ ╚═════╝ 
    
    🏺 翁氏宗谱RAG系统 v1.0.0
    📚 基于LangChain和深度求索模型的家谱知识问答系统
    """
    print(banner)

def run_tests():
    """运行基础功能测试"""
    print("🚀 开始运行基础功能测试...")
    print("=" * 50)
    
    success_count = 0
    total_tests = 0
    
    # 测试1：配置模块
    print("\n📋 测试1：配置管理模块")
    try:
        from src.config import config
        config.validate_config()
        print("✅ 配置模块测试通过")
        success_count += 1
    except Exception as e:
        print(f"❌ 配置模块测试失败: {str(e)}")
    total_tests += 1
    
    # 测试2：PDF处理模块
    print("\n📄 测试2：PDF处理模块")
    try:
        from src.document_processor import test_pdf_processor
        if test_pdf_processor():
            success_count += 1
    except Exception as e:
        print(f"❌ PDF处理模块测试失败: {str(e)}")
    total_tests += 1
    
    # 测试3：LLM连接
    print("\n🤖 测试3：LLM连接测试")
    try:
        from src.generator import test_llm_connection
        if test_llm_connection():
            success_count += 1
    except Exception as e:
        print(f"❌ LLM连接测试失败: {str(e)}")
    total_tests += 1
    
    # 显示测试结果
    print("\n" + "=" * 50)
    print(f"📊 测试完成！通过: {success_count}/{total_tests}")
    
    if success_count == total_tests:
        print("🎉 所有测试都通过了！系统状态良好。")
        return True
    elif success_count >= total_tests * 0.7:
        print("⚠️ 大部分测试通过，系统基本可用。")
        return True
    else:
        print("❌ 测试失败较多，请检查系统配置。")
        return False

def run_stage2_tests():
    """运行第二阶段功能测试"""
    print("🚀 开始运行第二阶段功能测试...")
    print("=" * 50)
    
    success_count = 0
    total_tests = 0
    
    # 测试1：向量化模块
    print("\n🔢 测试1：文本向量化模块")
    try:
        from src.embeddings import test_embedding_generator
        if test_embedding_generator():
            success_count += 1
    except Exception as e:
        print(f"❌ 向量化模块测试失败: {str(e)}")
    total_tests += 1
    
    # 测试2：检索模块
    print("\n🔍 测试2：语义检索模块")
    try:
        from src.retriever import test_retrieval_system
        if test_retrieval_system():
            success_count += 1
    except Exception as e:
        print(f"❌ 检索模块测试失败: {str(e)}")
    total_tests += 1
    
    # 测试3：RAG引擎
    print("\n⚙️ 测试3：RAG核心引擎")
    try:
        from src.rag_engine import test_rag_engine
        if test_rag_engine():
            success_count += 1
    except Exception as e:
        print(f"❌ RAG引擎测试失败: {str(e)}")
    total_tests += 1
    
    # 显示测试结果
    print("\n" + "=" * 50)
    print(f"📊 第二阶段测试完成！通过: {success_count}/{total_tests}")
    
    if success_count == total_tests:
        print("🎉 所有第二阶段测试都通过了！RAG功能已就绪。")
        return True
    elif success_count >= total_tests * 0.7:
        print("⚠️ 大部分第二阶段测试通过，RAG功能基本可用。")
        return True
    else:
        print("❌ 第二阶段测试失败较多，请检查RAG组件配置。")
        return False

def run_initialization():
    """运行系统初始化"""
    print("🔧 开始系统初始化...")
    print("=" * 50)
    
    try:
        # 导入RAG引擎
        from src.rag_engine import RAGEngine
        
        # 创建RAG引擎
        rag_engine = RAGEngine()
        print("✅ RAG引擎创建成功")
        
        # 从PDF初始化
        print("📄 正在从PDF文件初始化系统...")
        init_result = rag_engine.initialize_from_pdf()
        
        if init_result['success']:
            print("✅ 系统初始化成功！")
            print(f"   📊 处理统计:")
            stats = init_result['statistics']
            print(f"      - 原始文本长度: {stats['original_text_length']} 字符")
            print(f"      - 清洗后文本长度: {stats['cleaned_text_length']} 字符")
            print(f"      - 文本块数量: {stats['text_chunks_count']} 个")
            print(f"      - 处理耗时: {stats['processing_time']:.2f} 秒")
            print(f"      - LLM服务可用: {'是' if stats['llm_available'] else '否'}")
            
            # 显示PDF信息
            pdf_info = init_result.get('pdf_info', {})
            if 'error' not in pdf_info:
                print(f"   📋 PDF信息:")
                print(f"      - 页数: {pdf_info.get('page_count', '未知')}")
                print(f"      - 文件大小: {pdf_info.get('file_size', 0)} 字节")
            
            return True
        else:
            print(f"❌ 系统初始化失败: {init_result.get('message', '未知错误')}")
            return False
        
    except Exception as e:
        print(f"❌ 系统初始化失败: {str(e)}")
        return False

def run_qa_mode():
    """运行问答模式"""
    print("💬 启动RAG问答模式...")
    print("=" * 50)
    print("💡 提示：输入 'quit' 或 'exit' 退出")
    print()
    
    try:
        # 导入RAG引擎
        from src.rag_engine import RAGEngine, QAProcessor
        
        # 创建RAG引擎和问答处理器
        rag_engine = RAGEngine()
        qa_processor = QAProcessor(rag_engine)
        
        # 检查系统状态
        status = rag_engine.get_system_status()
        print(f"🔧 系统状态:")
        print(f"   - RAG引擎已初始化: {'是' if status['rag_engine_initialized'] else '否'}")
        print(f"   - 文档已处理: {'是' if status['document_processed'] else '否'}")
        print(f"   - LLM服务可用: {'是' if status['llm_available'] else '否'}")
        
        if not status['document_processed']:
            print("\n⚠️ 文档尚未处理，正在初始化...")
            init_result = rag_engine.initialize_from_pdf()
            if not init_result['success']:
                print(f"❌ 初始化失败: {init_result.get('message', '未知错误')}")
                print("🔄 将使用简单模式（无检索增强）")
        
        print("\n🚀 RAG问答系统已就绪！")
        print()
        
        while True:
            try:
                question = input("\n❓ 请输入您的问题: ").strip()
                
                if question.lower() in ['quit', 'exit', '退出']:
                    print("👋 再见！")
                    break
                
                if not question:
                    continue
                
                print("🤔 正在思考...")
                
                # 使用问答处理器处理问题
                result = qa_processor.process_question(question)
                
                if result['success']:
                    print(f"\n🤖 回答: {result['answer']}")
                    print(f"📊 模式: {result.get('mode', '未知')}")
                    
                    if result.get('mode') == 'rag':
                        retrieval_results = result.get('retrieval_results', [])
                        if retrieval_results:
                            print(f"🔍 找到 {len(retrieval_results)} 个相关文档片段")
                            print(f"📏 上下文长度: {result.get('context_length', 0)} 字符")
                else:
                    print(f"❌ 生成回答失败: {result.get('message', '未知错误')}")
                
            except KeyboardInterrupt:
                print("\n👋 再见！")
                break
            except Exception as e:
                print(f"❌ 处理问题时出错: {str(e)}")
        
        return True
        
    except Exception as e:
        print(f"❌ 问答模式启动失败: {str(e)}")
        return False

def run_web_mode():
    """运行Web界面模式"""
    print("🌐 启动Web界面...")
    print("💡 Web界面功能开发中，敬请期待！")
    return False

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='翁氏宗谱RAG系统')
    parser.add_argument('--mode', '-m', 
                       choices=['test', 'init', 'qa', 'web', 'stage2'], 
                       default='test',
                       help='运行模式 (test: 测试, init: 初始化, qa: 问答, web: Web界面, stage2: 第二阶段测试)')
    
    args = parser.parse_args()
    
    # 打印横幅
    print_banner()
    
    # 显示当前配置
    print("⚙️ 当前配置:")
    print(f"   ollama服务器: {config.get_ollama_url()}")
    print(f"   使用模型: {config.OLLAMA_MODEL}")
    print(f"   PDF文件: {config.PDF_PATH}")
    print("=" * 60)
    
    # 根据模式运行相应功能
    success = False
    
    if args.mode == 'test':
        success = run_tests()
    elif args.mode == 'stage2':
        success = run_stage2_tests()
    elif args.mode == 'init':
        success = run_initialization()
    elif args.mode == 'qa':
        success = run_qa_mode()
    elif args.mode == 'web':
        success = run_web_mode()
    else:
        print(f"❌ 未知的运行模式: {args.mode}")
        parser.print_help()

    # 退出
    if success:
        sys.exit(0)
    else:
        sys.exit(1)

if __name__ == "__main__":
    main() 